#include "../Diagnostic/Diagnostic1D.h"
#include "../Species/Species.h"
#include "../Parallel/ParallelMPI1D.h"
#include "../PhysicalField/PhysicalField.h"
#include "../Collision/Collision.h"
#include "../Field/Field1D.h"
#include "../Species/Species1D.h"

#include <iomanip>
#include <fstream>

using namespace std;

Diagnostic1D::Diagnostic1D(PicParams* params_, ParallelMPI* pmpi_, Grid* grid_, PhysicalField* fields_, vector<Species*>* vecSpecies_, vector<Collision*>* vecCollision_, vector<PSI*>* vecPSI_) :
    Diagnostic(params_, pmpi_, grid_, fields_, vecSpecies_, vecCollision_, vecPSI_)
{
	ParallelMPI1D* pmpi1D = static_cast<ParallelMPI1D*>(pmpi);
	index_domain_begin = pmpi1D->getCellStartingGlobalIndex(0);

    dx  = params->cell_length[0];
    dx_inv_  = 1.0 / params->cell_length[0];

	n_energy = 100;
	energy_max = 200.0;

    double B_magnitude = pow(params->externB[0], 2) + pow(params->externB[1], 2) + pow(params->externB[2], 2);
    B_magnitude = sqrt(B_magnitude);
    sin_alpha = params->externB[0] / B_magnitude;
    cos_alpha = sqrt(1.0 - sin_alpha * sin_alpha);

    ptclNum1D = new Field1D<double>(params->dims_field, "ptclNum");

	particle_number.resize(n_species);
	total_particle_energy.resize(n_species);

	particle_flux_left.resize(n_species);
	particle_flux_right.resize(n_species);
	heat_flux_left.resize(n_species);
	heat_flux_right.resize(n_species);

	angle_distribution_left.resize(n_species);
	angle_distribution_right.resize(n_species);
	energy_distribution_left.resize(n_species);
	energy_distribution_right.resize(n_species);

	for(int ispec = 0; ispec < n_species; ispec++)
	{
		angle_distribution_left[ispec].resize(90);
		angle_distribution_right[ispec].resize(90);
	}

	for(int ispec = 0; ispec < n_species; ispec++)
	{
		energy_distribution_left[ispec].resize(n_energy);
		energy_distribution_right[ispec].resize(n_energy);
	}

	psi_rate_left.resize(n_psi);
	psi_rate_right.resize(n_psi);
    psi_species1.resize(n_psi);
	for(int i_psi = 0; i_psi < n_psi; i_psi++)
	{
        //psi_species1[i_psi] = vecPSI[i_psi]->species1;
	}

	radiative_energy_collision.resize(n_collision);
	for(int i_collision = 0; i_collision < n_collision; i_collision++)
	{
        radiative_energy_collision[i_collision].resize(params->dims_space_global[0]);
	}

	sigma_left.resize(n_species);
	sigma_right.resize(n_species);
	for(int ispec = 0; ispec < n_species; ispec++)
	{
		sigma_left[ispec]  = 0.0;
		sigma_right[ispec] = 0.0;
	}

    //=============== vdf edf ===================================
    if(params->diagnostic_param.size() == 1)
    {
        init_vedf();
    }


}

Diagnostic1D::~Diagnostic1D()
{
	delete ptclNum1D;
}


void Diagnostic1D::run( ParallelMPI* pmpi, Grid* grid, vector<Species*>& vecSpecies, PhysicalField* fields, vector<PSI*>& vecPSI, int itime )
{
    Species1D *species1D;
    double v_square, v_magnitude, energy;
    double mass_ov_2, pi_ov_2;
    double wlt;			//weight * cell_length / time for calculating flux
    int i_angle, i_energy;
    double flux_temp;
    vector<double> angle_temp;
    vector<double> sigma_temp;

    angle_temp.resize(90);
    sigma_temp.resize(n_species);

    pi_ov_2 = params->const_pi * 0.5;

    //reset diagnostic parameters to zero
    if( ((itime - 1) % params->timestep_io) == 0 )
    {
        for(int ispec = 0; ispec < n_species; ispec++)
        {
            particle_flux_left[ispec] = 0.0;
            particle_flux_right[ispec] = 0.0;
            heat_flux_left[ispec] = 0.0;
            heat_flux_right[ispec] = 0.0;
            for(int i_angle = 0; i_angle < 90; i_angle++)
            {
                angle_distribution_left[ispec][i_angle] = 0.0;
                angle_distribution_right[ispec][i_angle] = 0.0;
            }
            for(int i_energy = 0; i_energy < n_energy; i_energy++)
            {
                energy_distribution_left[ispec][i_energy] = 0.0;
                energy_distribution_right[ispec][i_energy] = 0.0;
            }
        }

        for(int i_psi = 0; i_psi < n_psi; i_psi++)
        {
            psi_rate_left[i_psi] = 0.0;
            psi_rate_right[i_psi] = 0.0;
        }

        for(int i_collision = 0; i_collision < n_collision; i_collision++)
        {
            for(int i_bin = 0; i_bin < radiative_energy_collision[i_collision].size(); i_bin++)
            {
                radiative_energy_collision[i_collision][i_bin] = 0.0;
            }
        }
    }

    //calculate particle flux, heat flux, angle distribution, energy distribution
    if( (itime % params->timestep_io) > (params->timestep_io - params->timestep_average) || (itime % params->timestep_io) == 0 )
    {
        for(int ispec = 0; ispec < n_species; ispec++)
        {
            species1D = static_cast<Species1D*>(vecSpecies[ispec]);
            auto& psi_particles = species1D->psi_particles_;
            mass_ov_2 = 0.5 * species1D->species_param.mass;

            for(auto &particle : psi_particles)
            {
                //cout<<"particle number:  "<<p1->position(0,i_particle)<<endl;
                v_square = particle.velocity[0] * particle.velocity[0] + particle.velocity[1] * particle.velocity[1] + particle.velocity[2] * particle.velocity[2];
                v_magnitude = sqrt(v_square);
                energy = mass_ov_2 * v_square;
                i_angle = 90.0 * acos( abs(particle.velocity[0]) / v_magnitude ) / pi_ov_2;
                i_energy = n_energy * (energy / params->const_e) / energy_max;
                if(i_energy >= n_energy)
                {
                    i_energy = n_energy - 1;
                }
                if(particle.position[0] < 0.0 )
                {
                    particle_flux_left[ispec] += 1.0;
                    heat_flux_left[ispec] += energy;
                    if (i_angle >= 0 && i_angle < 90)
                    {
                        angle_distribution_left[ispec][i_angle] += 1.0;
                    }
                    else
                    {
                        log_warning<<"i_angle left out of range: i_angle = "<<i_angle;
                    }
                    energy_distribution_left[ispec][i_energy] += 1.0;

                }
                else if(particle.position[0] > params->sim_length[0])
                {
                    particle_flux_right[ispec] += 1.0;
                    heat_flux_right[ispec] += energy;
                    if (i_angle >= 0 && i_angle < 90)
                    {
                        angle_distribution_right[ispec][i_angle] += 1.0;
                    }
                    else
                    {
                        log_warning<<"i_angle right out of range: i_angle = "<<i_angle;
                    }
                    energy_distribution_right[ispec][i_energy] +=1.0;
                }
            }
        }
    }


    //MPI gather diagnostic parameters to master process
    if( (itime % params->timestep_io) == 0 )
    {
        vector<double> flux_temp;
        vector<double> angle_distribution_temp;
        vector<double> energy_distribution_temp;
        vector<double> psi_rate_temp;
        vector<double> radiative_energy_collision_temp;

        flux_temp.resize(n_species);
        angle_distribution_temp.resize(90);
        energy_distribution_temp.resize(n_energy);
        psi_rate_temp.resize(n_psi);
        radiative_energy_collision_temp.resize(params->dims_space_global[0]);

        //calculate average psi rate
        for(int i_psi = 0; i_psi < n_psi; i_psi++)
        {
            if(particle_flux_left[psi_species1[i_psi]] != 0.0)
            {
                psi_rate_left[i_psi]  /= particle_flux_left[psi_species1[i_psi]];
            }
            if(particle_flux_right[psi_species1[i_psi]] != 0.0)
            {
                psi_rate_right[i_psi] /= particle_flux_right[psi_species1[i_psi]];
            }
        }

        //same weight for all species
        species1D = static_cast<Species1D*>(vecSpecies[0]);
        wlt = species1D->species_param.weight / (params->cell_length_inv[0] * params->dt * params->timestep_average);
        for(int ispec = 0; ispec < n_species; ispec++)
        {
            particle_flux_left[ispec] *= wlt;
            particle_flux_right[ispec] *= wlt;
            heat_flux_left[ispec] *= wlt;
            heat_flux_right[ispec] *= wlt;
        }

        wlt = species1D->species_param.weight * params->const_e / (params->dt * params->timestep_average);
        for(int i_collision = 0; i_collision < n_collision; i_collision++)
        {
            for(int i_bin = 0; i_bin < radiative_energy_collision[i_collision].size(); i_bin++)
            {
                radiative_energy_collision[i_collision][i_bin] *= wlt;
            }
        }

        pmpi->reduce_sum_double(&particle_flux_left[0], &flux_temp[0], n_species);
        particle_flux_left = flux_temp;

        pmpi->reduce_sum_double(&particle_flux_right[0], &flux_temp[0], n_species);
        particle_flux_right = flux_temp;

        pmpi->reduce_sum_double(&heat_flux_left[0], &flux_temp[0], n_species);
        heat_flux_left = flux_temp;

        pmpi->reduce_sum_double(&heat_flux_right[0], &flux_temp[0], n_species);
        heat_flux_right = flux_temp;

        /*
        if(itime > params->timestep_io && pmpi->isMaster())
        {
            cout<<particle_flux_right[0]/wlt<<"  "<<particle_flux_right[1]/wlt<<endl;
        }
        */

        for(int ispec = 0; ispec < n_species; ispec++)
        {
            pmpi->reduce_sum_double(&angle_distribution_left[ispec][0], &angle_distribution_temp[0], 90);
            angle_distribution_left[ispec] = angle_distribution_temp;

            pmpi->reduce_sum_double(&angle_distribution_right[ispec][0], &angle_distribution_temp[0], 90);
            angle_distribution_right[ispec] = angle_distribution_temp;

            pmpi->reduce_sum_double(&energy_distribution_left[ispec][0], &energy_distribution_temp[0], n_energy);
            energy_distribution_left[ispec] = energy_distribution_temp;

            pmpi->reduce_sum_double(&energy_distribution_right[ispec][0], &energy_distribution_temp[0], n_energy);
            energy_distribution_right[ispec] = energy_distribution_temp;
        }

        pmpi->reduce_sum_double(&psi_rate_left[0], &psi_rate_temp[0], n_psi);
        psi_rate_left = psi_rate_temp;

        pmpi->reduce_sum_double(&psi_rate_right[0], &psi_rate_temp[0], n_psi);
        psi_rate_right = psi_rate_temp;

        for(int i_collision = 0; i_collision < n_collision; i_collision++)
        {
            pmpi->reduce_sum_double(&radiative_energy_collision[i_collision][0], &radiative_energy_collision_temp[0], radiative_energy_collision[i_collision].size());
            radiative_energy_collision[i_collision] = radiative_energy_collision_temp;
        }
    }


    //calculate sigma_left and sigma_right every timestep
    for(int ispec = 0; ispec < n_species; ispec++)
    {
        species1D = static_cast<Species1D*>(vecSpecies[ispec]);
        auto& psi_particles = species1D->psi_particles_;

        double sigma_left_temp, sigma_right_temp;
        sigma_left_temp = 0.0;
        sigma_right_temp = 0.0;

        wlt = species1D->species_param.charge * species1D->species_param.weight * params->cell_length[0];

        for(auto &particle : psi_particles)
        {
            if(particle.position[0] < 0.0)
            {
                sigma_left_temp += 1.0;
            }
            else if(particle.position[0] > params->sim_length[0])
            {
                sigma_right_temp += 1.0;
            }
        }

        sigma_left[ispec] += sigma_left_temp * wlt;
        sigma_right[ispec] += sigma_right_temp * wlt;
    }
    pmpi->reduce_sum_double(&sigma_left[0], &sigma_temp[0], n_species);
    sigma_left = sigma_temp;

    pmpi->reduce_sum_double(&sigma_right[0], &sigma_temp[0], n_species);
    sigma_right = sigma_temp;

    cal_VT(pmpi, vecSpecies, fields, itime);

    if(params->diagnostic_param.size() == 1)
    {
        cal_vedf(pmpi, vecSpecies, itime);
    }
    
}

//calculate velocity, temperature, total particle energy for each species, total electric field energy
//calculate every timestep
void Diagnostic1D::cal_VT(ParallelMPI* pmpi, vector<Species*>& vecSpecies, PhysicalField* fields, int itime)
{
    int i_temp;
    double xjn,xjmxi;
    double m_ov_3e;
    double m_ov_2;
    double vx, vy, vz;
    double timestep_average_inv_;
    double v_square;
    double total_electric_field_energy_temp;
    vector<int> particle_number_temp;
    vector<double> total_particle_energy_temp;

    particle_number_temp.resize(n_species);
    total_particle_energy_temp.resize(n_species);
    timestep_average_inv_ = 1.0 / params->timestep_average;
    for(int ispec = 0; ispec < vecSpecies.size(); ispec++)
    {
        //cout<<"cal_VT00000 "<<ispec<<endl;
        auto species1D = static_cast<Species1D*>(vecSpecies[ispec]);
        Field1D<double>* Vx1D_s = static_cast<Field1D<double>*>(fields->Vx_s[ispec]);
        Field1D<double>* Vy1D_s = static_cast<Field1D<double>*>(fields->Vy_s[ispec]);
        Field1D<double>* Vz1D_s = static_cast<Field1D<double>*>(fields->Vz_s[ispec]);
        Field1D<double>* Vp1D_s = static_cast<Field1D<double>*>(fields->Vp_s[ispec]);

        Field1D<double>* Vx1D_s_avg = static_cast<Field1D<double>*>(fields->Vx_s_avg[ispec]);
        Field1D<double>* Vy1D_s_avg = static_cast<Field1D<double>*>(fields->Vy_s_avg[ispec]);
        Field1D<double>* Vz1D_s_avg = static_cast<Field1D<double>*>(fields->Vz_s_avg[ispec]);
        Field1D<double>* Vp1D_s_avg = static_cast<Field1D<double>*>(fields->Vp_s_avg[ispec]);

        Field1D<double>* T1D_s = static_cast<Field1D<double>*>(fields->T_s[ispec]);
        Field1D<double>* T1D_s_avg = static_cast<Field1D<double>*>(fields->T_s_avg[ispec]);

        Field1D<double>* Tx1D_s = static_cast<Field1D<double>*>(fields->Tx_s[ispec]);
        Field1D<double>* Tx1D_s_avg = static_cast<Field1D<double>*>(fields->Tx_s_avg[ispec]);

        Field1D<double>* Ty1D_s = static_cast<Field1D<double>*>(fields->Ty_s[ispec]);
        Field1D<double>* Ty1D_s_avg = static_cast<Field1D<double>*>(fields->Ty_s_avg[ispec]);

        Field1D<double>* Tz1D_s = static_cast<Field1D<double>*>(fields->Tz_s[ispec]);
        Field1D<double>* Tz1D_s_avg = static_cast<Field1D<double>*>(fields->Tz_s_avg[ispec]);

        Field1D<double>* q1D_s = static_cast<Field1D<double>*>(fields->q_s[ispec]);
        Field1D<double>* q1D_s_avg = static_cast<Field1D<double>*>(fields->q_s_avg[ispec]);

        if(((itime - 1) % params->timestep_io) == 0)
        {
            Vx1D_s_avg->set_value(0.0);
            Vy1D_s_avg->set_value(0.0);
            Vz1D_s_avg->set_value(0.0);
            Vp1D_s_avg->set_value(0.0);
            T1D_s_avg ->set_value(0.0);
            q1D_s_avg ->set_value(0.0);

            if(params->is_temperature_anisotropic)
            {
                T1D_s_avg ->set_value(0.0);
                T1D_s_avg ->set_value(0.0);
                T1D_s_avg ->set_value(0.0);
            }

        }

        //calculate macroscopic velocity (average velocity) and particle number at grid points
        if( (itime % params->timestep_io) > (params->timestep_io - params->timestep_average) || (itime % params->timestep_io) == 0 )
        {
            m_ov_3e = species1D->species_param.mass / ( params->const_e * 3.0 );
            m_ov_2 = species1D->species_param.mass / 2.0;
            ptclNum1D->set_value(0.0);
            Vx1D_s->set_value(0.0);
            Vy1D_s->set_value(0.0);
            Vz1D_s->set_value(0.0);
            Vp1D_s->set_value(0.0);
            T1D_s->set_value(0.0);
            q1D_s->set_value(0.0);

            //reset particle_number and total_particle_energy
            particle_number[ispec] = 0;
            total_particle_energy[ispec] = 0.0;

            //cout<<"cal_VT1111 "<<endl;
            for(auto cell_particles : species1D->particles_)
            {
                for(auto& particle : cell_particles->data_)
                {
                    //Locate particle on the grid
                    xjn    = particle.position[0] * dx_inv_;  	//normalized distance to the first node
                    i_temp = floor(xjn);                   				//index of the central node
                    xjmxi  = xjn - (double)i_temp;              		//normalized distance to the nearest grid point

                    i_temp -= index_domain_begin;

                    double v_p = particle.velocity[0] * sin_alpha + particle.velocity[1] * cos_alpha;
                    (*ptclNum1D)(i_temp) 	+= 1.0;
                    (*Vx1D_s)(i_temp) 		+= particle.velocity[0];
                    (*Vy1D_s)(i_temp) 		+= particle.velocity[1];
                    (*Vz1D_s)(i_temp) 		+= particle.velocity[2];
                    (*Vp1D_s)(i_temp) 		+= v_p;

                    //calculate total total_particle_energy
                    v_square = particle.velocity[0] * particle.velocity[0] + particle.velocity[1] * particle.velocity[1] + particle.velocity[2] * particle.velocity[2];
                    total_particle_energy[ispec] += ( m_ov_2 * v_square );

                    (*q1D_s)(i_temp) 		+= species1D->species_param.mass * v_square * v_p;
                }
                particle_number[ispec] += cell_particles->size();
            }
            for(int i = 0; i < ptclNum1D->dims_[0]; i++)
            {
                if( (*ptclNum1D)(i) != 0.0 )
                {
                    (*Vx1D_s)(i) /= (*ptclNum1D)(i);
                    (*Vy1D_s)(i) /= (*ptclNum1D)(i);
                    (*Vz1D_s)(i) /= (*ptclNum1D)(i);
                    (*Vp1D_s)(i) /= (*ptclNum1D)(i);
                    (*q1D_s)(i) /= (*ptclNum1D)(i);
                }
            }

            //cout<<"cal_VT2222 "<<endl;


            //calculate temperature
            for(auto cell_particles : species1D->particles_)
            {
                for(auto& particle : cell_particles->data_)
                {
                    //Locate particle on the grid
                    xjn    = particle.position[0] * dx_inv_;  	//normalized distance to the first node
                    i_temp      = floor(xjn);                   		//index of the central node
                    xjmxi  = xjn - (double)i_temp;              		//normalized distance to the nearest grid point

                    i_temp -= index_domain_begin;

                    if(i_temp < 0 || i_temp >= Vx1D_s->dims_[0]) { cout<<"i_temp error: "<<i_temp<<"  "<<particle.position[0]<<endl; }
                    vx = particle.velocity[0] - (*Vx1D_s)(i_temp);
                    vy = particle.velocity[1] - (*Vy1D_s)(i_temp);
                    vz = particle.velocity[2] - (*Vz1D_s)(i_temp);
                    //cout<<T1D_s->dims_[0]<<"  "<<Vx1D_s->dims_[0]<<endl;
                    (*T1D_s)(i_temp) += ( vx * vx + vy * vy + vz * vz );

                    if(params->is_temperature_anisotropic)
                    {
                        (*Tx1D_s)(i_temp) += 3.0 * (vx * vx);
                        (*Ty1D_s)(i_temp) += 3.0 * (vy * vy);
                        (*Tz1D_s)(i_temp) += 3.0 * (vz * vz);
                    }
                }
            }


            for(int i = 0; i < ptclNum1D->dims_[0]; i++)
            {
                if( (*ptclNum1D)(i) != 0.0 )
                {
                    (*T1D_s)(i) = (*T1D_s)(i) * m_ov_3e / (*ptclNum1D)(i);

                    if(params->is_temperature_anisotropic)
                    {
                        (*Tx1D_s)(i) = (*Tx1D_s)(i) * m_ov_3e / (*ptclNum1D)(i);
                        (*Ty1D_s)(i) = (*Ty1D_s)(i) * m_ov_3e / (*ptclNum1D)(i);
                        (*Tz1D_s)(i) = (*Tz1D_s)(i) * m_ov_3e / (*ptclNum1D)(i);
                    }
                }

            }

            //cout<<"cal_VT3333 "<<endl;

            //sum velocity and temperature
            for(int i = 0; i < ptclNum1D->dims_[0]; i++)
            {
                (*Vx1D_s_avg)(i) += (*Vx1D_s)(i);
                (*Vy1D_s_avg)(i) += (*Vy1D_s)(i);
                (*Vz1D_s_avg)(i) += (*Vz1D_s)(i);
                (*Vp1D_s_avg)(i) += (*Vp1D_s)(i);
                (*T1D_s_avg)(i)  += (*T1D_s)(i);
                (*q1D_s_avg)(i)  += (*q1D_s)(i);

                if(params->is_temperature_anisotropic)
                {
                    (*Tx1D_s_avg)(i)  += (*Tx1D_s)(i);
                    (*Ty1D_s_avg)(i)  += (*Ty1D_s)(i);
                    (*Tz1D_s_avg)(i)  += (*Tz1D_s)(i);
                }
            }

            //cout<<"cal_VT4444 "<<endl;
        }


        //Calculate the average parameters and MPI gather
        if( (itime % params->timestep_io) == 0 )
        {
            for(int i = 0; i < ptclNum1D->dims_[0]; i++)
            {
                (*Vx1D_s_avg)(i) *= timestep_average_inv_;
                (*Vy1D_s_avg)(i) *= timestep_average_inv_;
                (*Vz1D_s_avg)(i) *= timestep_average_inv_;
                (*Vp1D_s_avg)(i) *= timestep_average_inv_;
                (*T1D_s_avg)(i)  *= timestep_average_inv_;
                (*q1D_s_avg)(i)  *= timestep_average_inv_;

                if(params->is_temperature_anisotropic)
                {
                    (*Tx1D_s_avg)(i)  *= timestep_average_inv_;
                    (*Ty1D_s_avg)(i)  *= timestep_average_inv_;
                    (*Tz1D_s_avg)(i)  *= timestep_average_inv_;
                }
            }

            if(itime == 0)
            {
                for(int i = 0; i < ptclNum1D->dims_[0]; i++)
                {
                    (*Vx1D_s_avg)(i) /= timestep_average_inv_;
                    (*Vy1D_s_avg)(i) /= timestep_average_inv_;
                    (*Vz1D_s_avg)(i) /= timestep_average_inv_;
                    (*Vp1D_s_avg)(i) /= timestep_average_inv_;
                    (*T1D_s_avg)(i)  /= timestep_average_inv_;
                    (*q1D_s_avg)(i)  /= timestep_average_inv_;

                    if(params->is_temperature_anisotropic)
                    {
                        (*Tx1D_s_avg)(i)  /= timestep_average_inv_;
                        (*Ty1D_s_avg)(i)  /= timestep_average_inv_;
                        (*Tz1D_s_avg)(i)  /= timestep_average_inv_;
                    }
                }
            }

            //another way: firstly gather V, T, ptclNum, then calculate V_global, T_global
            ParallelMPI1D* pmpi1D = static_cast<ParallelMPI1D*>(pmpi);
            pmpi1D->gatherField2( fields->Vx_s_global_avg[ispec], Vx1D_s_avg );
            pmpi1D->gatherField2( fields->Vy_s_global_avg[ispec], Vy1D_s_avg );
            pmpi1D->gatherField2( fields->Vz_s_global_avg[ispec], Vz1D_s_avg );
            pmpi1D->gatherField2( fields->Vp_s_global_avg[ispec], Vp1D_s_avg );
            pmpi1D->gatherField2( fields->T_s_global_avg [ispec], T1D_s_avg );
            pmpi1D->gatherField2( fields->q_s_global_avg [ispec], q1D_s_avg );

            if(params->is_temperature_anisotropic)
            {
                pmpi1D->gatherField2( fields->Tx_s_global_avg [ispec], Tx1D_s_avg );
                pmpi1D->gatherField2( fields->Ty_s_global_avg [ispec], Ty1D_s_avg );
                pmpi1D->gatherField2( fields->Tz_s_global_avg [ispec], Tz1D_s_avg );
            }

        }


    }


    //sum particle_number and total_particle_energy to master process
    if( (itime % params->timestep_io) == 0 )
    {
        pmpi->reduce_sum_int(&particle_number[0], &particle_number_temp[0], n_species);
        particle_number = particle_number_temp;
        pmpi->reduce_sum_double(&total_particle_energy[0], &total_particle_energy_temp[0], n_species);
        total_particle_energy = total_particle_energy_temp;
    }


    //calculate electric field energy
    Field1D<double>* Ex1D = static_cast<Field1D<double>*>(fields->Ex_);
    for(int i = 0; i < Ex1D->dims_[0]; i++)
    {
        total_electric_field_energy += (*Ex1D)(i) * (*Ex1D)(i);
    }

    pmpi->reduce_sum_double(&total_electric_field_energy, &total_electric_field_energy_temp, 1);
    total_electric_field_energy = total_electric_field_energy_temp;

}


//init vedf
void Diagnostic1D::init_vedf()
{
    vector<int> vedf_dims(3);
    vedf_dims[0] = diagnostic_param.vedf_n_position;
    vedf_dims[1] = diagnostic_param.vedf_n_species;
    vedf_dims[2] = diagnostic_param.vedf_n_velocity;

    vdf_parallel_to_B.allocate_dims(vedf_dims);
    vdf_perpendicular_to_B.allocate_dims(vedf_dims);
    vdf_z.allocate_dims(vedf_dims);
    edf.allocate_dims(vedf_dims);
    vedf_temp.allocate_dims((vedf_dims));

    for(int i_position = 0; i_position < diagnostic_param.vedf_n_position; i_position++)
    {
        double position = diagnostic_param.vedf_position[i_position];
        for(int i_cell = 0; i_cell < params->dims_space[0]; i_cell++)
        {
            double cell_start = params->local_min[0] + i_cell * params->cell_length[0];
            double cell_end   = cell_start + params->cell_length[0];
            if(position >= cell_start - 1.0e-10 && position < cell_end - 1.0e-10)
            {
                vdf_cell_local.push_back(i_cell);
                vdf_position_index_global.push_back(i_position);
                log_debug<<"vdf cell number: "<<i_cell<<" "<<i_position;
            }
        }
    }

    vedf_dims.resize(2);
    vedf_dims[0] = diagnostic_param.vedf_n_species;
    vedf_dims[1] = diagnostic_param.vedf_n_velocity;
    vdf_velocity_array.allocate_dims(vedf_dims);
    edf_energy_array.allocate_dims(vedf_dims);
    vdf_dv.resize(diagnostic_param.vedf_n_species);
    edf_de.resize(diagnostic_param.vedf_n_species);

    for(int i_species = 0; i_species < diagnostic_param.vedf_n_species; i_species++)
    {
        int species_number = diagnostic_param.vedf_species_number_list[i_species];
        double v_max = sqrt(2.0 * diagnostic_param.edf_max_energy[i_species] * params->const_e / params->species_param[species_number].mass);
        double v_min = -v_max;

        vdf_dv[i_species] = (v_max - v_min) / diagnostic_param.vedf_n_velocity;
        for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
        {
            if(i_velocity == 0)
            {
                vdf_velocity_array(i_species, i_velocity) = v_min + 0.5 * vdf_dv[i_species];
            }
            else
            {
                vdf_velocity_array(i_species, i_velocity) = vdf_velocity_array(i_species, i_velocity-1) + vdf_dv[i_species];
            }
            //cout<<"vdf_velocity_array(i_species, i_velocity)  "<<vdf_velocity_array(i_species, i_velocity)<<endl;
        }

        edf_de[i_species] = diagnostic_param.edf_max_energy[i_species] / diagnostic_param.vedf_n_velocity;
        for(int i_energy = 0; i_energy < diagnostic_param.vedf_n_velocity; i_energy++)
        {
            if(i_energy == 0)
            {
                edf_energy_array(i_species, i_energy) = 0.5 * edf_de[i_species];
            }
            else
            {
                edf_energy_array(i_species, i_energy) = edf_energy_array(i_species, i_energy-1) + edf_de[i_species];
            }
        }
    }
}

//calculate velocity/energy distribution function
void Diagnostic1D::cal_vedf(ParallelMPI* pmpi, vector<Species*>& vecSpecies, int itime)
{

    //reset diagnostic parameters to zero
    if( ((itime - 1) % params->timestep_io) == 0 )
    {
        vdf_parallel_to_B.set_value(0.0);
        vdf_perpendicular_to_B.set_value(0.0);
        vdf_z.set_value(0.0);
        edf.set_value(0.0);
    }

    //calculate in timestep_average
    if( (itime % params->timestep_io) > (params->timestep_io - params->timestep_average) || (itime % params->timestep_io) == 0 )
    {
        for(int i = 0; i < vdf_cell_local.size(); i++)
        {
            int i_cell_local = vdf_cell_local[i];
            int i_position_index_global = vdf_position_index_global[i];
            for(int i_species = 0; i_species < diagnostic_param.vedf_n_species; i_species++)
            {
                auto species1D = static_cast<Species1D*>(vecSpecies[diagnostic_param.vedf_species_number_list[i_species]]);

                for(int i_extend = -1; i_extend < 2; i_extend++)
                {
                    auto cell_particles = species1D->particles_[i_cell_local+i_extend];
                    for(auto &particle : cell_particles->data_)
                    {
                        double v_max = sqrt(2.0 * diagnostic_param.edf_max_energy[i_species] * params->const_e / species1D->species_param.mass);
                        double v_min = -v_max;
                        double v_d = (v_max - v_min) / diagnostic_param.vedf_n_velocity;
                        int vdf_n_velocity_half = diagnostic_param.vedf_n_velocity / 2;
                        int i_v;

                        //vdf parallel to B
                        i_v = (particle.velocity[0] * sin_alpha + particle.velocity[1] * cos_alpha) / v_d + vdf_n_velocity_half;
                        if(i_v < 0)
                        {
                            i_v = 0;
                        }
                        else if(i_v >= diagnostic_param.vedf_n_velocity)
                        {
                            i_v = diagnostic_param.vedf_n_velocity - 1;
                        }
                        vdf_parallel_to_B(i_position_index_global, i_species, i_v) += 1.0;

                        //vdf perpendicular to B in the xy plane
                        i_v = (particle.velocity[0] * cos_alpha - particle.velocity[1] * sin_alpha) / v_d + vdf_n_velocity_half;
                        if(i_v < 0)
                        {
                            i_v = 0;
                        }
                        else if(i_v >= diagnostic_param.vedf_n_velocity)
                        {
                            i_v = diagnostic_param.vedf_n_velocity - 1;
                        }
                        vdf_perpendicular_to_B(i_position_index_global, i_species, i_v) += 1.0;

                        //vdf in z direction
                        i_v = particle.velocity[2] / v_d + vdf_n_velocity_half;
                        if(i_v < 0)
                        {
                            i_v = 0;
                        }
                        else if(i_v >= diagnostic_param.vedf_n_velocity)
                        {
                            i_v = diagnostic_param.vedf_n_velocity - 1;
                        }
                        vdf_z(i_position_index_global, i_species, i_v) += 1.0;


                        //edf
                        double particle_energy = 0.5 * species1D->species_param.mass * (pow(particle.velocity[0], 2) + pow(particle.velocity[1], 2) + pow(particle.velocity[2], 2));
                        particle_energy /= params->const_e;
                        int i_e = diagnostic_param.vedf_n_velocity * particle_energy / diagnostic_param.edf_max_energy[i_species];
                        //cout<<"particle energy "<<particle_energy<<"  "<<i_e<<endl;

                        if(i_e < 0)
                        {
                            i_e = 0;
                        }
                        else if(i_e >= diagnostic_param.vedf_n_velocity)
                        {
                            i_e = diagnostic_param.vedf_n_velocity - 1;
                        }
                        edf(i_position_index_global, i_species, i_e) += 1.0;
                    }
                }

            }
        }
    }


    //MPI gather diagnostic parameters to master process
    if( (itime % params->timestep_io) == 0 )
    {
        pmpi->reduce_sum_double(vdf_parallel_to_B.data_, vedf_temp.data_, vdf_parallel_to_B.dims_global_);
        vdf_parallel_to_B = vedf_temp;

        pmpi->reduce_sum_double(vdf_perpendicular_to_B.data_, vedf_temp.data_, vdf_perpendicular_to_B.dims_global_);
        vdf_perpendicular_to_B = vedf_temp;

        pmpi->reduce_sum_double(vdf_z.data_, vedf_temp.data_, vdf_z.dims_global_);
        vdf_z = vedf_temp;

        pmpi->reduce_sum_double(edf.data_, vedf_temp.data_, edf.dims_global_);
        edf = vedf_temp;

        //calculate vdf: f_v = vdf / (tot_number * dv)
        for(int i_position = 0; i_position < diagnostic_param.vedf_n_position; i_position++)
        {
            for(int i_species = 0; i_species < diagnostic_param.vedf_n_species; i_species++)
            {
                double particle_number_total;

                particle_number_total = 0.0;
                for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
                {
                    particle_number_total += vdf_parallel_to_B(i_position, i_species, i_velocity);
                }
                for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
                {
                    vdf_parallel_to_B(i_position, i_species, i_velocity) /= (particle_number_total * vdf_dv[i_species]);
                }

                particle_number_total = 0.0;
                for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
                {
                    particle_number_total += vdf_perpendicular_to_B(i_position, i_species, i_velocity);
                }
                for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
                {
                    vdf_perpendicular_to_B(i_position, i_species, i_velocity) /= (particle_number_total * vdf_dv[i_species]);
                }

                particle_number_total = 0.0;
                for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
                {
                    particle_number_total += vdf_z(i_position, i_species, i_velocity);
                }
                for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
                {
                    vdf_z(i_position, i_species, i_velocity) /= (particle_number_total * vdf_dv[i_species]);
                }

                particle_number_total = 0.0;
                for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
                {
                    particle_number_total += edf(i_position, i_species, i_velocity);
                }
                for(int i_velocity = 0; i_velocity < diagnostic_param.vedf_n_velocity; i_velocity++)
                {
                    edf(i_position, i_species, i_velocity) /= (particle_number_total * edf_de[i_species]);
                }
            }
        }

    }


}
